﻿using Orleans;
using Orleans.Providers;
using Orleans.Runtime;
using Orleans.Storage;
using System.Threading.Tasks;

namespace iTool.ClusterComponent
{
    public interface IKeyValueStorageService : iToolServiceWithStringKey
    {
        Task<KeyValueState> GetStateAsync();
        Task<int> GetVersionAsync();
        Task<string> GetAsync();
        Task SetAsync(string bytes);
        Task RemoveAsync();
    }

    [StorageProvider(ProviderName = "KeyValueStorageService")]
    public class KeyValueStorageService : Grain<KeyValueState>, IKeyValueStorageService, IGrainStorage
    {
        readonly IKeyValueStorageProvider storageProvider;
        public KeyValueStorageService(IKeyValueStorageProvider storageProvider) 
        {
            this.storageProvider = storageProvider;
        }

        public Task<string> GetAsync()
        {
            return Task.FromResult(this.State.Value);
        }
        public async Task SetAsync(string bytes)
        {
            this.State.Version++;
            this.State.Value = bytes;
            await this.WriteStateAsync();
        }

        public async Task RemoveAsync()
        {
            await this.ClearStateAsync();
            this.DeactivateOnIdle();
        }

        public async Task ClearStateAsync(string grainType, GrainReference grainReference, IGrainState grainState)
        {
            var key = this.GetKey(grainType, grainReference);
            await this.storageProvider.RemoveAsync(key);
        }

        public async Task ReadStateAsync(string grainType, GrainReference grainReference, IGrainState grainState)
        {
            var key = this.GetKey(grainType, grainReference);
            var type = grainState.State.GetType();
            grainState.State = await this.storageProvider.GetOrAddAsync(key, () => new KeyValueState(), type) ?? new KeyValueState();
        }

        public async Task WriteStateAsync(string grainType, GrainReference grainReference, IGrainState grainState)
        {
            var key = this.GetKey(grainType, grainReference);
            await this.storageProvider.AddOrUpdateAsync(key, grainState.State, grainState.State.GetType());
        }

        private string GetKey(string className, GrainReference grainReference)
        {
            string key = grainReference.GetPrimaryKeyString();
            return $"{className}_{key}";
        }

        public Task<KeyValueState> GetStateAsync()
        {
            return Task.FromResult(this.State);
        }

        public Task<int> GetVersionAsync()
        {
            return Task.FromResult(this.State.Version);
        }
    }


    public class KeyValueState 
    {
        public int Version { get; set; }
        public string Value { get; set; }
    }
}
